package com.hanter.android.radlib.crypto;

import android.annotation.SuppressLint;
import android.content.Context;
import android.os.Build;
import android.os.Environment;

import com.hanter.android.radlib.util.StringUtils;

import java.io.File;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.IOException;
import java.security.SecureRandom;
import java.security.spec.KeySpec;

import javax.crypto.Cipher;
import javax.crypto.KeyGenerator;
import javax.crypto.SecretKeyFactory;
import javax.crypto.spec.IvParameterSpec;
import javax.crypto.spec.PBEKeySpec;
import javax.crypto.spec.SecretKeySpec;

/**
 * 参考：
 * https://shoewann0402.github.io/2016/06/10/security-crypto-provider-deprecated-in-android-n/
 * https://android.googlesource.com/platform/development/+/master/samples/BrokenKeyDerivation
 */
public class AesUtils {

    private static final String AES_TRANSFORMATION = "AES/CBC/PKCS5PADDING";
    private static final String AES_SALT_NAME = "aes_salt";

    private static File baseDir = Environment.getDataDirectory();

    public static void init(Context context) {
        baseDir = context.getFilesDir();
    }

    public static String encrypt(String seed, String cleartext)
            throws Exception {
        byte[] rawKey = getRawKeySecurely(seed);
        byte[] result = encrypt(rawKey, cleartext.getBytes(StandardCharsets.UTF_8));
        return StringUtils.bytesToHex(result);
    }

    public static String encryptInsecurely(String seed, String cleartext)
            throws Exception {
        byte[] rawKey = getRawKeyInsecurely(seed);
        byte[] result = encrypt(rawKey, cleartext.getBytes(StandardCharsets.UTF_8));
        return StringUtils.bytesToHex(result);
    }

    public static String decrypt(String seed, String encrypted)
            throws Exception {
        byte[] rawKey = getRawKeySecurely(seed);
        byte[] enc = StringUtils.hexToBytes(encrypted);
        byte[] result = decrypt(rawKey, enc);
        return new String(result);
    }

    public static String decryptInsecurely(String seed, String encrypted)
            throws Exception {
        byte[] rawKey = getRawKeyInsecurely(seed);
        byte[] enc = StringUtils.hexToBytes(encrypted);
        byte[] result = decrypt(rawKey, enc);
        return new String(result);
    }

    private static byte[] getRawKeyInsecurely(String seed) throws Exception {
        byte[] seedBytes = seed.getBytes(StandardCharsets.UTF_8);
        KeyGenerator keyGen = KeyGenerator.getInstance("AES");
        if (Build.VERSION.SDK_INT >= Build.VERSION_CODES.N) {
            return InsecureSHA1PRNGKeyDerivator.deriveInsecureKey(seedBytes, seedBytes.length);
        } else {
            @SuppressLint("DeletedProvider")
            SecureRandom sr = SecureRandom.getInstance("SHA1PRNG", "Crypto");
            sr.setSeed(seedBytes);
            keyGen.init(128, sr); // 192 and 256 bits may not be available
            return keyGen.generateKey().getEncoded();
        }
    }

    private static byte[] getRawKeySecurely(String password) throws Exception {
        int iterationCount = 50;
        int keyLength = 128; // 256-bits for AES-256, 128-bits for AES-128, etc
        int saltLength = 16; // bytes; should be the same size as the output (256 / 8 = 32)
        byte[] salt = retrieveSalt(saltLength);

        /* Use this to derive the key from the password: */
        KeySpec keySpec = new PBEKeySpec(password.toCharArray(), salt,
                iterationCount, keyLength);
        SecretKeyFactory keyFactory = SecretKeyFactory.getInstance("PBKDF2withHmacSHA1");
        return keyFactory.generateSecret(keySpec).getEncoded();
    }

    private static byte[] retrieveSalt(int saltLength) {
        // Salt must be at least the same size as the key.
        byte[] salt = new byte[saltLength];
        // Create a random salt if encrypting for the first time, and save it for future use.
        readFromFileOrCreateRandom(AES_SALT_NAME, salt);
        return salt;
    }

    private static byte[] encrypt(byte[] raw, byte[] clear) throws Exception {
        SecretKeySpec skSpec = new SecretKeySpec(raw, "AES");
        Cipher cipher = Cipher.getInstance(AES_TRANSFORMATION);
        cipher.init(Cipher.ENCRYPT_MODE, skSpec, new IvParameterSpec(
                new byte[cipher.getBlockSize()]));
        return cipher.doFinal(clear);
    }

    private static byte[] decrypt(byte[] raw, byte[] encrypted)
            throws Exception {
        SecretKeySpec skSpec = new SecretKeySpec(raw, "AES");
        Cipher cipher = Cipher.getInstance(AES_TRANSFORMATION);
        cipher.init(Cipher.DECRYPT_MODE, skSpec, new IvParameterSpec(
                new byte[cipher.getBlockSize()]));
        return cipher.doFinal(encrypted);
    }

    /**
     * Read from file or return random bytes in the given array.
     *
     * <p>Save to file if file didn't exist.
     */
    private static void readFromFileOrCreateRandom(String fileName, byte[] bytes) {
        if (fileExists(fileName)) {
            readBytesFromFile(fileName, bytes);
            return;
        }
        SecureRandom sr = new SecureRandom();
        sr.nextBytes(bytes);
        writeToFile(fileName, bytes);
    }

    private static boolean fileExists(String fileName) {
        File file = new File(baseDir, fileName);
        return file.exists();
    }

    private static boolean removeFile(String fileName) {
        File file = new File(baseDir, fileName);
        return file.delete();
    }

    private static void writeToFile(String fileName, byte[] bytes) {
        FileOutputStream fos = null;
        try {
            fos = new FileOutputStream(new File(baseDir, fileName));
            fos.write(bytes);
        } catch (IOException e) {
            throw new RuntimeException("Couldn't write to " + fileName, e);
        } finally {
            if (fos != null) {
                try {
                    fos.close();
                } catch (IOException e) {
                    e.printStackTrace();
                }
            }
        }
    }

    private static void readBytesFromFile(String fileName, byte[] bytes) {
        FileInputStream fis = null;
        try {
            fis = new FileInputStream(new File(baseDir, fileName));
            int numBytes = 0;
            while (numBytes < bytes.length) {
                int n = fis.read(bytes, numBytes, bytes.length - numBytes);
                if (n <= 0) {
                    throw new RuntimeException("Couldn't read from " + fileName);
                }
                numBytes += n;
            }
        } catch (IOException e) {
            throw new RuntimeException("Couldn't read from " + fileName, e);
        } finally {
            if (fis != null) {
                try {
                    fis.close();
                } catch (IOException e) {
                    e.printStackTrace();
                }
            }
        }
    }

}
